package org.jberney.wvdexporter.video;

import java.awt.image.BufferedImage;

public class WvdFrame {

	private BufferedImage image;

	public WvdFrame(WvdFrameHeader wvdFrameHeader, byte[] frameBytes) {

		byte[][] yPlane = new byte[wvdFrameHeader.getWidth()][wvdFrameHeader
				.getHeight()];
		byte[][] uPlane = new byte[wvdFrameHeader.getWidth() / 2][wvdFrameHeader
				.getHeight() / 2];
		byte[][] vPlane = new byte[wvdFrameHeader.getWidth() / 2][wvdFrameHeader
				.getHeight() / 2];

		for (int y = 0; y < wvdFrameHeader.getHeight() / 2; y++) {

			for (int x = 0; x < wvdFrameHeader.getWidth() / 2; x++) {

				int srcOffset = 2 * (y * wvdFrameHeader.getWidth() + 2 * x);

				yPlane[x][y] = frameBytes[srcOffset];
				yPlane[x + wvdFrameHeader.getWidth() / 2][y] = frameBytes[srcOffset + 1];
				uPlane[x][y] = frameBytes[srcOffset + 2];
				vPlane[x][y] = frameBytes[srcOffset + 3];

			}

		}

		int offsetConstant = wvdFrameHeader.getWidth()
				* wvdFrameHeader.getHeight();

		for (int y = 0; y < wvdFrameHeader.getHeight() / 4; y++) {

			for (int x = 0; x < wvdFrameHeader.getWidth() / 2; x++) {

				int srcOffset = offsetConstant + 2
						* (y * wvdFrameHeader.getWidth() + 2 * x);

				yPlane[x][y + wvdFrameHeader.getHeight() / 2] = frameBytes[srcOffset];
				yPlane[x + wvdFrameHeader.getWidth() / 2][y
						+ wvdFrameHeader.getHeight() / 2] = frameBytes[srcOffset + 1];
				yPlane[x][y + 3 * wvdFrameHeader.getHeight() / 4] = frameBytes[srcOffset + 2];
				yPlane[x + wvdFrameHeader.getWidth() / 2][y + 3
						* wvdFrameHeader.getHeight() / 4] = frameBytes[srcOffset + 3];

			}

		}

		this.image = new BufferedImage(wvdFrameHeader.getWidth(),
				wvdFrameHeader.getHeight(), BufferedImage.TYPE_3BYTE_BGR);

		for (int j = 0; j < wvdFrameHeader.getHeight(); j++) {

			for (int i = 0; i < wvdFrameHeader.getWidth(); i++) {

				int y = yPlane[i][j] & 0xFF;
				int u = uPlane[i / 2][j / 2] & 0xFF;
				int v = vPlane[i / 2][j / 2] & 0xFF;

				int r = this.clamp(y + 1.4020 * (v - 128));
				int g = this.clamp(y - 0.3441 * (v - 128) - 0.7141 * (u - 128));
				int b = this.clamp(y + 1.7720 * (u - 128));

				int rgb = (0xff << 24) | (r << 16) | (g << 8) | b;

				this.image.setRGB(i, j, rgb);

			}

		}

	}

	private int clamp(double x) {

		int clamped = (int) x;

		if (x < 0.0) {
			clamped = 0;
		} else if (x > 255.0) {
			clamped = 255;
		}

		return clamped;
	}

	public BufferedImage getFrameAsImage() {
		return this.image;
	}

}
